# -*- coding: utf-8 -*-
import os
import imageio
import numpy as np
import pickle

# 解压缩，返回解压后的字典
def unpickle(file):
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='latin1')
    return dict

# 创建文件夹结构
def create_class_directories(base_path, num_classes=10):
    for i in range(num_classes):
        os.makedirs(os.path.join(base_path, f'class_{i}'), exist_ok=True)

# 生成训练集图片
create_class_directories('train')  # 创建训练集类别文件夹
for j in range(1, 6):
    dataName = f"data_batch_{j}"
    Xtr = unpickle(dataName)
    print(dataName + " is loading...")

    for i in range(10000):
        img = np.reshape(Xtr['data'][i], (3, 32, 32))
        img = img.transpose(1, 2, 0)  # 读取image
        label = Xtr['labels'][i]
        # 构建类别文件夹和文件名
        picName = f'train/class_{label}/{label}_{i + (j - 1)*10000}.jpg'
        imageio.imsave(picName, img)
    print(dataName + " loaded.")

# 生成测试集图片
create_class_directories('test')  # 创建测试集类别文件夹
print("test_batch is loading...")
testXtr = unpickle("test_batch")
for i in range(10000):
    img = np.reshape(testXtr['data'][i], (3, 32, 32))
    img = img.transpose(1, 2, 0)
    label = testXtr['labels'][i]
    picName = f'test/class_{label}/{label}_{i}.jpg'
    imageio.imsave(picName, img)
print("test_batch loaded.")
